from typing import Tuple
from dataclasses import dataclass
import pandas as pd
import numpy as np

from . import base


@dataclass
class Equal(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'equal'
    col_a: str = 'close'
    col_b: str = 'close'
    labels: Tuple[int, int] = (1, 0)

    @classmethod
    def trans(cls, a: pd.Series, b: pd.Series, labels: Tuple[int, int]):
        y, n = labels
        return np.where(a == b, y, n)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col_a], data[self.col_b], self.labels)
        return data


@dataclass
class Eval(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'eval'
    expr: str = ''

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = data.eval(self.expr)
        return data


@dataclass
class Div(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'div'
    col_num: str = 'close'
    col_deno: str = 'close'

    @classmethod
    def trans(cls, a: pd.Series, b: pd.Series):
        return a / b

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col_num], data[self.col_deno])
        return data
